from src.gen import main
from src.enums import DocumentSubset

# https://open.toutiao.com/a7243055925005763106
import os
os.environ["OPENAI_API_KEY"] = "Your openai key"
from langchain import OpenAI, PromptTemplate, LLMChain
from langchain.text_splitter import CharacterTextSplitter
from langchain.chains.mapreduce import MapReduceChain
from langchain.prompts import PromptTemplate

llm = OpenAI(temperature=0)
text_splitter = CharacterTextSplitter()

with open("./state_of_the_union.txt") as f:
    state_of_the_union = f.read()
docs = text_splitter.split_text(state_of_the_union)[:4]

from langchain.chains.summarize import load_summarize_chain
import textwrap

chain = load_summarize_chain(llm, chain_type="map_reduce", verbose = True)

output_summary = chain.run(docs)
wrapped_text = textwrap.fill(output_summary, width=100)
print(wrapped_text)

prompt_template = """请用简洁的要点来总结以下内容： {text}

简洁的要点总结："""
BULLET_POINT_PROMPT = PromptTemplate(template=prompt_template, input_variables=["text"])

chain = load_summarize_chain(llm, 
                             chain_type="stuff", 
                             prompt=BULLET_POINT_PROMPT)

output_summary = chain.run(docs)

wrapped_text = textwrap.fill(output_summary, 
                             width=100,
                             break_long_words=False,
                             replace_whitespace=False)
print(wrapped_text)

chain = load_summarize_chain(llm, chain_type="refine")

output_summary = chain.run(docs)
wrapped_text = textwrap.fill(output_summary, width=100)
print(wrapped_text)
